package com.dianmi.seckill.common.concurrent;

import com.dianmi.seckill.common.advice.ResponseCode;
import com.dianmi.seckill.common.advice.ServiceException;
import com.dianmi.seckill.common.utils.JsonUtil;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;
import java.util.function.BiFunction;

/**
 * 两个参数的多任务执行服务
 * @author xieqiaofu
 * @date 2020/12/2
 */
@Slf4j
public class BiMultiTaskWaitResultService<T, V, R> {

    private Executor executor;

    private ThreadLocal<CountDownLatch> latchThreadLocal = new ThreadLocal<>();

    private ThreadLocal<ConcurrentHashMap<Integer, R>> resultTheadLocal = new ThreadLocal<ConcurrentHashMap<Integer, R>>();

    private ThreadLocal<Integer> index = new ThreadLocal<>();

    private ThreadLocal<Integer> taskCount = new ThreadLocal<>();

    /**
     * 队列大小和线程池线程数需要业务方自行评估
     */
    public BiMultiTaskWaitResultService(Executor executor) {
        this.executor = executor;
    }

    /**
     * 任务数量
     * @param taskNum
     */
    public void start(int taskNum) {
        CountDownLatch latch = new CountDownLatch(taskNum);
        latchThreadLocal.set(latch);

        ConcurrentHashMap<Integer, R> map = new ConcurrentHashMap<>();

        resultTheadLocal.set(map);

        index.set(0);

        taskCount.set(taskNum);
    }

    public void submit(BiFunction<T, V, R> function, T t, V v) {
        Integer idx = index.get();
        CountDownLatch latch = latchThreadLocal.get();
        ConcurrentHashMap<Integer, R> map = resultTheadLocal.get();

        if (idx >= taskCount.get()) {
            throw new ServiceException(ResponseCode.QUANTITY_OVERFLOW);
        }
        executor.execute(new Runnable() {
            @Override
            public void run() {
                try {
                    R r = function.apply(t, v);
                    map.put(idx, r);
                } catch (Exception e) {
                    log.error("任务执行异常," + JsonUtil.toJson(t), e);
                } finally {
                    latch.countDown();
                }
            }
        });
        index.set(idx+1);
    }

    /**
     * 一定要加超时时间，不然会阻塞。超时时间长短业务方自行评估
     * @param time
     * @param timeUnit
     * @return
     */
    public List<R> get(long time, TimeUnit timeUnit) {
        CountDownLatch latch = latchThreadLocal.get();
        try {
            latch.await(time, timeUnit);
            ConcurrentHashMap<Integer, R> resultMap = resultTheadLocal.get();
            List<R> ret = new ArrayList<>();
            for (int i=0, len=index.get(); i<len; i++) {
                ret.add(resultMap.get(i));
            }
            return ret;
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        } finally {
            clear();
        }
    }

    private void clear() {
        latchThreadLocal.remove();
        resultTheadLocal.remove();
        index.remove();
        taskCount.remove();
    }
}
